{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 500, Loss: 0.39196738600730896\n",
      "Epoch 1000, Loss: 0.38889023661613464\n",
      "Epoch 1500, Loss: 0.3881550133228302\n",
      "Epoch 2000, Loss: 0.388113796710968\n",
      "Epoch 2500, Loss: 0.3880850076675415\n",
      "Epoch 3000, Loss: 0.3880627751350403\n",
      "Epoch 3500, Loss: 0.388047993183136\n",
      "Epoch 4000, Loss: 0.3880394399166107\n",
      "Epoch 4500, Loss: 0.38803350925445557\n",
      "Epoch 5000, Loss: 0.3880285620689392\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "element 0 of tensors does not require grad and does not have a grad_fn",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_15960\\3037395295.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     50\u001b[0m \u001b[0mN_pred\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mt_test\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     51\u001b[0m \u001b[0mx_pred\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mx0\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mt_test\u001b[0m \u001b[1;33m*\u001b[0m \u001b[0mN_pred\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 52\u001b[1;33m \u001b[0mu_pred\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx_pred\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mt_test\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgrad_outputs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mones_like\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mt_test\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mx_pred\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     53\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     54\u001b[0m \u001b[1;31m# Convert to numpy for plotting\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Conda\\lib\\site-packages\\torch\\autograd\\__init__.py\u001b[0m in \u001b[0;36mgrad\u001b[1;34m(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched, materialize_grads)\u001b[0m\n\u001b[0;32m    392\u001b[0m         )\n\u001b[0;32m    393\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 394\u001b[1;33m         result = Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[0;32m    395\u001b[0m             \u001b[0mt_outputs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    396\u001b[0m             \u001b[0mgrad_outputs_\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: element 0 of tensors does not require grad and does not have a grad_fn"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# Define the neural network\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.hidden1 = nn.Linear(1, 20)\n",
    "        self.hidden2 = nn.Linear(20, 20)\n",
    "        self.output = nn.Linear(20, 1)\n",
    "\n",
    "    def forward(self, t):\n",
    "        t = torch.tanh(self.hidden1(t))\n",
    "        t = torch.tanh(self.hidden2(t))\n",
    "        t = self.output(t)\n",
    "        return t\n",
    "\n",
    "# Initialize the neural network\n",
    "net = Net()\n",
    "\n",
    "# Define the initial condition\n",
    "x0 = torch.tensor([1.0])\n",
    "\n",
    "# Define the loss function\n",
    "def loss_function(t):\n",
    "    t = t.view(-1, 1)\n",
    "    t.requires_grad = True\n",
    "    N = net(t)\n",
    "    x = x0 + t * N\n",
    "    u = torch.autograd.grad(x, t, grad_outputs=torch.ones_like(t), create_graph=True)[0] + x\n",
    "    cost = x**2 + u**2\n",
    "    return torch.mean(cost)\n",
    "\n",
    "# Training the neural network\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.001)\n",
    "epochs = 5000\n",
    "t_train = torch.linspace(0, 1, 100).view(-1, 1)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    optimizer.zero_grad()\n",
    "    loss = loss_function(t_train)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if (epoch + 1) % 500 == 0:\n",
    "        print(f'Epoch {epoch+1}, Loss: {loss.item()}')\n",
    "\n",
    "# Predict the state and control\n",
    "t_test = torch.linspace(0, 1, 100).view(-1, 1)\n",
    "N_pred = net(t_test).detach()\n",
    "x_pred = x0 + t_test * N_pred\n",
    "u_pred = torch.autograd.grad(x_pred, t_test, grad_outputs=torch.ones_like(t_test), create_graph=True)[0] + x_pred\n",
    "\n",
    "# Convert to numpy for plotting\n",
    "t_test_np = t_test.numpy()\n",
    "x_pred_np = x_pred.numpy()\n",
    "u_pred_np = u_pred.numpy()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
